package coordinator

import (
	"sync"
	"time"

	"vdrift/internal/logger"
	pb "vdrift/pkg/proto"
)

type Election struct {
	mu         sync.RWMutex
	peers      map[string]pb.NodeState
	localID    string
	quorumSize int
	leaderID   string
	leaderCh   chan string
}

func NewElection(localID string, peerCount int) *Election {
	return &Election{
		peers:      make(map[string]pb.NodeState),
		localID:    localID,
		quorumSize: (peerCount / 2) + 1,
		leaderCh:   make(chan string, 1),
	}
}

func (e *Election) Update(state pb.NodeState) {
	e.mu.Lock()
	e.peers[state.NodeId] = state
	e.mu.Unlock()
	e.recalculate()
}

func (e *Election) GetLeader() string {
	e.mu.RLock()
	defer e.mu.RUnlock()
	return e.leaderID
}

func (e *Election) recalculate() {
	e.mu.Lock()
	defer e.mu.Unlock()

	valid := make(map[string]pb.NodeState)
	for id, s := range e.peers {
		if time.Since(time.Unix(0, s.Timestamp)) < 5*time.Second {
			valid[id] = s
		}
	}

	if len(valid) < e.quorumSize {
		if e.leaderID != "" {
			logger.WithComponent("election").Warn("lost quorum, unbinding VIP")
		}
		e.leaderID = ""
		e.notifyLeader()
		return
	}

	var leader string
	maxP := int32(-1)
	for id, s := range valid {
		if s.Healthy && s.Priority > maxP {
			maxP = s.Priority
			leader = id
		}
	}

	if e.leaderID != leader {
		logger.WithElection(leader, maxP, e.quorumSize).Info("election result")
		e.leaderID = leader
		e.notifyLeader()
	}
}

func (e *Election) notifyLeader() {
	select {
	case e.leaderCh <- e.leaderID:
	default:
	}
}

func (e *Election) LeaderChan() <-chan string {
	return e.leaderCh
}
